#include <stdio.h>
#include <math.h>
#include <stdlib.h>

#define PI 3.1415926

float *fun(float fpeak, float dt, int nt) {
    float T, t, *ric;
    float temp1, temp2;
    int it;
    T = 1. / fpeak;
    ric = (float *) malloc(sizeof(float) * nt);
    for (it = 0; it < nt; it++)ric[it] = 0;
    for (it = 0; it < nt; it++) {
        t = it * dt - T;
        temp1 = PI * fpeak * t;
        temp2 = temp1 * temp1;
        ric[it] = (1 - 2 * temp2) * exp(-temp2);
    }
    return ric;
}

void updatep(float **p, float **vx, float **vz, float **v, float **den, float dx, float dz, float dt, int nx, int nz,
             float **data) {
    int ix, iz;
    float dx0, dz0;
    dx0 = dt / dx;
    dz0 = dt / dz;
    for (ix = 1; ix < nx - 1; ix++)
        for (iz = 1; iz < nz - 1; iz++) {
            p[ix][iz] += ((vx[ix][iz] - vx[ix - 1][iz]) * dx0 + (vz[ix][iz] - vz[ix][iz - 1]) * dz0) * den[ix][iz] *
                         v[ix][iz] * v[ix][iz] + data[ix][iz];
        }
}

void updatev(float **p, float **vx, float **vz, float **den, float dx, float dz, float dt, int nx, int nz) {
    int ix, iz;
    float dx0, dz0;
    dx0 = dt / dx;
    dz0 = dt / dz;
    for (ix = 1; ix < nx - 1; ix++)
        for (iz = 1; iz < nz - 1; iz++) {
            vx[ix][iz] += 1. / den[ix][iz] * (p[ix + 1][iz] - p[ix][iz]) * dx0;
            vz[ix][iz] += 1. / den[ix][iz] * (p[ix][iz + 1] - p[ix][iz]) * dz0;
        }
}


int main() {
    FILE *fpin, *fpout;
    char out[100], in[100];
    int sx0, sz0, sd0;
    float fpeak, dt, *ric, dz, dx, **den;
    float sx, sdepth, gdepth;
    float **data, **p, **vx, **vz, **v, **record;
    int it, nt, ix, iz, nx, nz;
    fpin = fopen("../par/wave.txt", "r");
    fscanf(fpin, "fpeak=%f\n", &fpeak);
    fscanf(fpin, "dt=%f\n", &dt);
    fscanf(fpin, "nt=%d\n", &nt);
    fscanf(fpin, "out=%s\n", out);
    fscanf(fpin, "nx=%d\n", &nx);
    fscanf(fpin, "nz=%d\n", &nz);
    fscanf(fpin, "dx=%f\n", &dx);
    fscanf(fpin, "dz=%f\n", &dz);
    fscanf(fpin, "sx=%f\n", &sx);
    fscanf(fpin, "sdepth=%f\n", &sdepth);
    fscanf(fpin, "gdepth=%f\n", &gdepth);
    fclose(fpin);
    ric = fun(fpeak, dt, nt);
    sx0 = (int) (sx / dx + 0.5);
    sz0 = (int) (sdepth / dz + 0.5);
    sd0 = (int) (gdepth / dz + 0.3);
    p = (float **) malloc(sizeof(float *) * nx);
    vx = (float **) malloc(sizeof(float *) * nx);
    vz = (float **) malloc(sizeof(float *) * nx);
    v = (float **) malloc(sizeof(float *) * nx);
    den = (float **) malloc(sizeof(float *) * nx);
    data = (float **) malloc(sizeof(float *) * nx);
    record = (float **) malloc(sizeof(float *) * nx);
    for (ix = 0; ix < nx; ix++) {
        record[ix] = (float *) malloc(sizeof(float) * nt);
        p[ix] = (float *) malloc(sizeof(float) * nz);
        vx[ix] = (float *) malloc(sizeof(float) * nz);
        vz[ix] = (float *) malloc(sizeof(float) * nz);
        data[ix] = (float *) malloc(sizeof(float) * nz);
        v[ix] = (float *) malloc(sizeof(float) * nz);
        den[ix] = (float *) malloc(sizeof(float) * nz);
        for (iz = 0; iz < nz; iz++) {
            p[ix][iz] = 0;
            vx[ix][iz] = 0;
            vz[ix][iz] = 0;
            data[ix][iz] = 0;
            den[ix][iz] = 1;
            v[ix][iz] = 1500;
        }
    }
    
    for(iz=0;iz<nz;iz++)
    {
        for(ix=0;ix<nx;ix++)
        {
            if(iz>500)
            v[ix][iz]=1500;
            else if(ix-iz>=500)
            v[ix][iz]=700;
        }
    }
    
    fpout = fopen("../data/wave_rick.txt", "w");
    for (it = 0; it < nt; it++)
        fprintf(fpout, "%f\t%f\n", it * dt, ric[it]);
    fclose(fpout);
    for (it = 0; it < nt; it++) {
        data[sx0][sd0] = ric[it];
        updatep(p, vx, vz, v, den, dx, dz, dt, nx, nz, data);
        updatev(p, vx, vz, den, dx, dz, dt, nx, nz);
        for (ix = 0; ix < nx; ix++)
            record[ix][it] = p[ix][sd0];
        if (it % 4 == 0)
            printf("%d\n", it);
    }
    fpout = fopen("../data/wave.bin", "wb");
    for (ix = 0; ix < nx; ix++)
        fwrite(p[ix], sizeof(float), nz, fpout);
    fclose(fpout);
    fpout = fopen("../data/record.bin", "wb");
    for (ix = 0; ix < nx; ix++)
        fwrite(p[ix], sizeof(float), nt, fpout);
    fclose(fpout);
    for (ix = 0; ix < nx; ix++) {
        free(v[ix]);
        free(vx[ix]);
        free(vz[ix]);
        free(data[ix]);
        free(den[ix]);
    }
    free(data);
    free(den);
    free(vx);
    free(vz);
    free(v);
    free(ric);
    return 0;
}
